package com.weather.websocket;

import org.springframework.web.socket.*;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.concurrent.ConcurrentHashMap;

public class WebSocketHandler extends TextWebSocketHandler {

    // 用户ID -> Session 的映射
    private static final ConcurrentHashMap<String, WebSocketSession> userSessionMap = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        // 从参数中获取 userId
        String userId = getUserIdFromSession(session);
        if (userId != null) {
            userSessionMap.put(userId, session);
            System.out.println("用户 [" + userId + "] 已连接");
        }
    }

    @Override
    public void handleTextMessage(WebSocketSession session, TextMessage message) {
        System.out.println("收到消息：" + message.getPayload());
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        String userId = getUserIdFromSession(session);
        if (userId != null) {
            userSessionMap.remove(userId);
            System.out.println("用户 [" + userId + "] 连接关闭");
        }
    }

    // 从 URI 中解析 userId 参数
    private String getUserIdFromSession(WebSocketSession session) {
        String uri = session.getUri().toString();
        if (uri.contains("userId=")) {
            return uri.substring(uri.indexOf("userId=") + 7);
        }
        return null;
    }

    // 公开一个方法供外部调用，用于服务端主动发送消息
    public static void sendMsgToUser(String userId, String msg) {
        WebSocketSession session = userSessionMap.get(userId);
        if (session != null && session.isOpen()) {
            try {
                session.sendMessage(new TextMessage(msg));
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}
